#!/bin/bash

VERSION=1.0
DEVICE="/dev/sr0"
LABEL="cyburn"
SPEED=6
ECCRED=0
CHECKSUMINVOKE=0
PROGSTATE=""
CHECKSUMBINARY=""
CHECKSUMNAME=""
MODE=""
DRM=""
OUTDIR=""

folder2iso () {
    echo "[cyburn] [$PROGSTATE] operating in $MODE mode, creating iso files from \"$1\" to \"$OUTDIR\""
    case $MODE in
        dvd) MAXBYTE=4700000000 ;;
        dvd9) MAXBYTE=8500000000 ;;
        bd) MAXBYTE=25000000000 ;;
        bd50) MAXBYTE=50000000000 ;;
        bdxl) MAXBYTE=100000000000 ;;
    esac

    if ! [[ $ECCRED -eq 0 ]]
    then
        printf "[cyburn] [$PROGSTATE] reducing $MODE target size to fill with ECC data ($MAXBYTE -> "
        ECCREDM=$(echo "$ECCRED*0.01" | bc)
        MAXBYTE=$(echo "$MAXBYTE-$MAXBYTE*$ECCREDM" | bc)
        MAXBYTE=${MAXBYTE%.*}
        printf "$MAXBYTE)\n"
    fi

    cd "$1"
    echo "[cyburn] [$PROGSTATE] preparing temporary directories..."
    DU=0
    DISK=1
    SKIP=0
    du -sb ./* | while read -r s f
    do
        if (( s >= MAXBYTE ))
        then
            echo "[cyburn] [$PROGSTATE] [WARNING] created d0 disk for larger files"
            mkdir -p "$OUTDIR"/d0
            SKIP=1
        elif (( DU + s >= MAXBYTE ))
        then
            (( DISK++ ))
            DU=0
        fi
        mkdir -p "$OUTDIR"/d$DISK
        if [[ $SKIP -eq 1 ]]
        then
            cp "$f" "$OUTDIR"/d0/
        else
            (( DU += s ))
            cp "$f" "$OUTDIR"/d$DISK/
        fi
        SKIP=0
    done
    cd "$OUTDIR"
    for d in ./d*
    do
        CLABEL="$LABEL""_$(basename -- $d)"
        if [ -d "$d" ]
        then
            echo "[cyburn] [$PROGSTATE] writing iso file for $CLABEL with mkisofs..."
            if ! mkisofs -v -Jr -preparer "cyburn v$VERSION" -V "$CLABEL" -o "$CLABEL.iso" "$d"
            then
                echo "[cyburn] [$PROGSTATE] [ERROR] an error occured while writing ISO file!"
                exit 1
            fi

            if ! [[ $ECCRED -eq 0 ]]
            then
                echo "[cyburn] [$PROGSTATE] filling ISO with dvdisaster RS03 ECC..."
                if ! dvdisaster -i "$CLABEL.iso" -mRS03 -c
                then
                    echo "[cyburn] [$PROGSTATE] failed to write dvdisaster RS03 ECC..."
                    exit 1
                fi
            fi
            
            if [[ $CHECKSUMINVOKE -eq 1 ]]
            then
                echo "[cyburn] [$PROGSTATE] calculating $CHECKSUMNAME checksum of $CLABEL.iso..."
                if ! $CHECKSUMBINARY "$CLABEL.iso" > "$CLABEL".iso.$CHECKSUMNAME
                then
                    echo "[cyburn] [$PROGSTATE] [ERROR] an error occured while calculating $CHECKSUMNAME checksum!"
                    exit 1
                fi
            fi

            echo "[cyburn] [$PROGSTATE] deleting temporary directory"
            rm -r "$d"
        fi
    done
    exit 0
}

iso2disk () {
    if [[ $CHECKSUMINVOKE -eq 1 ]]
    then
        echo "[cyburn] [$PROGSTATE] invoked $CHECKSUMNAME checksum integrity check"
        verify "$1"
    else
        echo "[cyburn] [$PROGSTATE] burning with "$SPEED"x speed"
        if ! growisofs -speed=$SPEED -dvd-compat -Z $DEVICE="$1"
        then
            echo "[cyburn] [ERROR] growisofs failed to burn the image! (faulty disk?)"
            exit 1
        fi
        exit 0
    fi
}

disk2copy () {
    echo "[cyburn] [$PROGSTATE] creating $MODE copy from $DEVICE to \"$OUTDIR\""
    cd "$OUTDIR"
    mkdir -p ".tmpcyburn"

    if [[ $LABEL == "cyburn" ]]
    then
        set_label "$(isoinfo -d -i $DEVICE | sed -n 's/Volume id: //p')"
        echo "[cyburn] [$PROGSTATE] using original ISO label: $LABEL"
    fi

    if [[ $DRM == "css" && $MODE == "iso" ]]
    then
        echo "[cyburn] [$PROGSTATE] using makemkvcon to descramble and copy disk to iso..."
        if ! makemkvcon backup --decrypt disc:0 "$OUTDIR/.tmpcyburn/$LABEL.iso"
        then
            echo "[cyburn] [$PROGSTATE] [ERROR] makemkvcon failed to create a decrypted copy!"
            exit 1
        fi
        if ! mv "$OUTDIR/.tmpcyburn/$LABEL.iso" "$OUTDIR/$LABEL.iso"
        then
            echo "[cyburn] [$PROGSTATE] [ERROR] makemkvcon failed to create a decrypted iso copy!"
            exit 1
        fi
    elif [[ $DRM == "css" ]]
    then
        echo "[cyburn] [$PROGSTATE] using dvdbackup to descramble and copy disk..."
        if ! dvdbackup -r a -v -p -i $DEVICE -n ".tmpcyburn" -M
        then
            echo "[cyburn] [$PROGSTATE] [ERROR] dvdbackup failed to create a descrambled copy! The DVD may be authored with structural copy protection like intentional bad blocks or duplicate titles (try using -m iso to use makemkvcon for descrambling)."
            exit 1 
        fi
    elif [[ $DRM == "aacs" ]]
    then
        echo "[cyburn] [$PROGSTATE] using makemkvcon to decrypt and copy disk..."
        if ! makemkvcon backup --decrypt disc:0 "$OUTDIR/.tmpcyburn"
        then
            echo "[cyburn] [$PROGSTATE] [ERROR] makemkvcon failed to create a decrypted copy!"
            exit 1
        fi
    elif [[ $DRM == "" && $MODE == "iso" ]]
    then
        echo "[cyburn] [$PROGSTATE] copying disk with dd..."
        if ! dd if=$DEVICE of="$OUTDIR/.tmpcyburn/$LABEL.iso"
        then
            echo "[cyburn] [$PROGSTATE] [ERROR] dd failed to copy disk! You may need root privileges."
            exit 1
        fi
        mv "$OUTDIR/.tmpcyburn/$LABEL.iso" "$OUTDIR/$LABEL.iso"
    fi

    if [[ $MODE == "tar" ]]
    then
        echo "[cyburn] [$PROGSTATE] setting read-only permissions for archiving..."
        find .tmpcyburn -type d -exec chmod 555 {} +
        find .tmpcyburn -type f -exec chmod 444 {} +
        echo "[cyburn] [$PROGSTATE] building tar file..."
        if ! tar -C "$OUTDIR/.tmpcyburn" -cf "$LABEL.tar" .
        then
            echo "[cyburn] [$PROGSTATE] [ERROR] an error occured while building the tar file!"
            exit 1
        fi
    elif [[ $DRM == "aacs" && $MODE == "iso" ]]
    then
        echo "[cyburn] [$PROGSTATE] building ISO file..."
        if ! mkisofs -v -preparer "cyburn v$VERSION" -V "$LABEL" -dvd-video -udf -o "$OUTDIR/$LABEL.iso" "$OUTDIR/.tmpcyburn"
        then
            echo "[cyburn] [$PROGSTATE] [ERROR] mkisofs failed to build the iso file!"
            exit 1
        fi
    elif [[ $MODE == "fs" ]]
    then
        mkdir -p "$LABEL"
        if ! mv .tmpcyburn/* "$LABEL"
        then
            echo "[cyburn] [$PROGSTATE] [ERROR] failed to move files to output directory!"
            exit 1
        fi
    fi

    if [[ $CHECKSUMINVOKE -eq 1 ]]
    then
        echo "[cyburn] [$PROGSTATE] calculating $CHECKSUMNAME checksum..."
        if [[ $MODE == "fs" ]]
        then
            if ! find "$LABEL" -type f -exec $CHECKSUMBINARY {} >> "$LABEL.$CHECKSUMNAME"  +
            then
                echo "[cyburn] [$PROGSTATE] [ERROR] an error occured while calculating $CHECKSUMNAME checksum!"
                exit 1
            fi
        else
            if ! $CHECKSUMBINARY "$LABEL.$MODE" > "$LABEL.$MODE.$CHECKSUMNAME"
            then
                echo "[cyburn] [$PROGSTATE] [ERROR] an error occured while calculating $CHECKSUMNAME checksum!"
                exit 1
            fi
        fi
    fi

    echo "[cyburn] [$PROGSTATE] deleting temporary directory"
    find .tmpcyburn -type d -exec chmod 755 {} +
    find .tmpcyburn -type f -exec chmod 644 {} +
    rm -r .tmpcyburn
    exit 0
}

verify () {
    if [[ ! -r "$1".$CHECKSUMNAME ]]
    then
        echo "[cyburn] [$PROGSTATE] $1.$CHECKSUMNAME checksum file not available, calculating checksum now..."
        IMAGESUM=$($CHECKSUMBINARY "$1" | cut -d " " -f1)
    else
        IMAGESUM=$(cat "$1".$CHECKSUMNAME | cut -d " " -f1)
    fi

    echo "[cyburn] [$PROGSTATE] ISO Checksum  = $IMAGESUM"

    IMAGECOUNT=$(expr $(du -b "$1" | awk '{print $1}') / 2048)
    DISKSUM=$(dd if=$DEVICE bs=2048 count=$IMAGECOUNT status=none | $CHECKSUMBINARY | cut -d " " -f1)

    echo "[cyburn] [$PROGSTATE] Disk Checksum = $DISKSUM"
    
    if [[ "$IMAGESUM" == "$DISKSUM" ]]
    then
        echo "[cyburn] [$PROGSTATE] Checksum integrity test passed!"
        exit 0
    else
        echo "[cyburn] [$PROGSTATE] Checksum integrity test failed!"
        exit 1
    fi
}

set_checksum_algo () {
    CHECKSUMBINARY=$1sum
    CHECKSUMNAME=$1
    CHECKSUMINVOKE=1
	if [ -x /usr/bin/$CHECKSUMBINARY ]
	then
		echo "[cyburn] using $CHECKSUMNAME (in /usr/bin/$CHECKSUMBINARY) as checksum algorithm"
	else
		echo "[cyburn] [ERROR] invalid checksum algorithm: $1"
		exit 1
	fi
}

set_output_directory () {
	if [ ! -d "$1" ]
	then
		mkdir -p "$1"
	fi
	if [ -z "$(ls -A "$1")" ]
	then
		OUTDIR="$1"
	else
		echo "[cyburn] [ERROR] invalid output directory (must be empty!): $1"
		exit 1
	fi
}

set_label () {
    if [[ ${#1} -gt 32 ]]
    then
        echo "[cyburn] [WARNING] invalid ISO label (${#1}/32 characters): $1"
    fi
    LABEL="$1"
}

print_help () {
	printf "\n\
folder2iso\n\
-f /path/to/folder                 Target directory\n\
-o /path/to/folder                 ISO Output directory\n\
-m [dvd,dvd9,bd,bd50,bdxl]         ISO target mode (default: dvd)\n\
-l \"cyburn\"                        ISO label\n\
-a [md5,sha1,sha256,sha512]        Create checksum file alongsite ISO\n\
-e [0-100]                         Fill at least e-percent of ISO with dvdisaster RS03 ECC\n\
\n\
iso2disk\n\
-d /dev/sr0                        Target device\n\
-b /path/to/image.iso              Burn image\n\
-s 6                               Burn speed\n\
-a [md5,sha1,sha256,sha512]        Verify burned image via checksum\n\
\n\
disk2copy\n\
-d /dev/sr0                        Target device\n\
-m [tar,iso,fs]                    Output mode (default: tar) (tar and fs mode only works with -x flag)\n\
-o /path/to/folder                 Output directory\n\
-l \"cyburn\"                        tar/iso/folder name (without suffix)\n\
-a [md5,sha1,sha256,sha512]        Calculate checksum of output\n\
-x [css,aacs]                      Bypass CSS (dvdbackup+libdvdcss for tar/fs and makemkvcon for iso with structural protection support), and AACS/AACS2/BD+ (makemkvcon) DRM copy protection\n"
	exit 0
}

echo "cyburn version $VERSION"
while getopts d:l:o:m:a:f:b:c:e:x:s:h OPTION
do
	case "${OPTION}"
	in
        d) DEVICE="${OPTARG}" ;;
        m) MODE="${OPTARG}" ;;
        l) set_label "${OPTARG}" ;;
        o) set_output_directory "$(readlink -f "${OPTARG}")" ;;
        a) set_checksum_algo "${OPTARG}" ;;
        s) SPEED="${OPTARG}" ;;
        f) PROGSTATE="folder2iso" && F2IA="$(readlink -f "${OPTARG}")" ;;
        b) PROGSTATE="iso2disk" && I2DA="$(readlink -f "${OPTARG}")" ;;
        e) ECCRED="${OPTARG}" ;;
        x) DRM="${OPTARG}" ;;
        h) print_help && exit 0;;
	esac
done

if [[ $PROGSTATE == "folder2iso" ]]
then
    if [[ $MODE == "" ]]
    then
        MODE="dvd"
    fi
    folder2iso "$F2IA"
elif [[ $PROGSTATE == "iso2disk" ]]
then
    iso2disk "$I2DA"
elif [[ $OUTDIR != "" ]]
then
    PROGSTATE="disk2copy"
    if [[ $MODE == "" ]]
    then
        MODE="tar"
    fi
    disk2copy
else
    print_help
fi